import argparse
import asyncio
import fcntl
import json
import os

import aiohttp


async def load_request_ids(REQUEST_ID_FILE_PATH):
    if os.path.exists(REQUEST_ID_FILE_PATH):
        with open(REQUEST_ID_FILE_PATH, 'r') as file:
            return json.load(file)
    return {}


async def save_request_ids(request_ids, REQUEST_ID_FILE_PATH):
    with open(REQUEST_ID_FILE_PATH, 'w') as file:
        fcntl.flock(file, fcntl.LOCK_EX)
        json.dump(request_ids, file, ensure_ascii=False, indent=4)
        fcntl.flock(file, fcntl.LOCK_UN)


async def re_request_ids(ip, request_id, REQUEST_ID_FILE_PATH, remove):
    request_ids = await load_request_ids(REQUEST_ID_FILE_PATH)
    request_ids.setdefault(ip, [])
    if remove:
        try:
            request_ids[ip].remove(request_id)
        except ValueError:
            pass
    else:
        if request_id not in request_ids[ip]:
            request_ids[ip].append(request_id)
    await save_request_ids(request_ids, REQUEST_ID_FILE_PATH)


async def stop_request(ip, session, request_id, model_name, port, REQUEST_ID_FILE_PATH):
    headers = {"Accept": "application/json", "Content-type": "application/json"}
    url = f'http://{ip}:{port}/v2/models/{model_name}/stopInfer'
    payload = {"id": request_id}

    try:
        async with session.post(url=url, json=payload, headers=headers) as resp:
            if resp.status == 200:
                print(f"\nIP {ip} 停止 {request_id} 请求成功: {resp.status}")
                await re_request_ids(ip, request_id, REQUEST_ID_FILE_PATH, True)
            else:
                print(f"\nIP {ip} 停止 {request_id} 请求错误: {resp.status}")
    except Exception as e:
        print(f"\nIP {ip} 停止 {request_id} 请求出错: {type(e).__name__} - {e}")


async def stop_requests_based_on_ids(ip, ids, model_name, port, REQUEST_ID_FILE_PATH):
    async with aiohttp.ClientSession() as session:
        tasks = []
        for id_ in ids:
            tasks.append(stop_request(ip, session, id_, model_name, port, REQUEST_ID_FILE_PATH))
        await asyncio.gather(*tasks)


def load_request_ids(REQUEST_ID_FILE_PATH):
    if os.path.exists(REQUEST_ID_FILE_PATH):
        with open(REQUEST_ID_FILE_PATH, 'r') as file:
            return json.load(file)
    return {}


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--port', type=int, default=1025)
    parser.add_argument('--model_name', type=str, default='deepseek-r1')
    parser.add_argument('--id_file', type=str, default='./src/request_ids.json')
    args = parser.parse_args()
    REQUEST_ID_FILE_PATH = args.id_file
    request_ids = load_request_ids(REQUEST_ID_FILE_PATH)
    for ip, ids in request_ids.items():
        asyncio.run(stop_requests_based_on_ids(ip, ids, args.model_name, args.port, REQUEST_ID_FILE_PATH))
